package com.dxb.test.mybatis;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.scripting.xmltags.DynamicContext;
import org.apache.ibatis.scripting.xmltags.DynamicSqlSource;
import org.apache.ibatis.scripting.xmltags.SqlNode;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import java.lang.reflect.Field;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.HashSet;
import java.util.Set;

@Intercepts({@Signature(method = "update", args = {MappedStatement.class, Object.class}, type = Executor.class),
        @Signature(method = "query", args = {MappedStatement.class, Object.class, RowBounds.class,
                ResultHandler.class}, type = Executor.class)})
public class SchemaInterceptor implements Interceptor {

    protected static Log LOG = LogFactory.getLog(SchemaInterceptor.class);
    private static final String SCHEMA = "schema";

    private String schema;

    private Set<Integer> sourceStorage = new HashSet<Integer>();

    @Override
    public Object intercept(Invocation invocation) throws Throwable {

        Object[] args = invocation.getArgs();

        MappedStatement mappedStatement = (MappedStatement) args[0];

        SqlSource sqlSource = mappedStatement.getSqlSource();

        // 只拦截动态sql
        if (sqlSource instanceof DynamicSqlSource) {

            // 获取到sqlNode对象
            Field field = DynamicSqlSource.class.getDeclaredField("rootSqlNode");
            field.setAccessible(true);
            SqlNode sqlnode = (SqlNode) field.get(sqlSource);
            // 获取动态代理对象
            SqlNode proxyNode = proxyNode(sqlnode);

            field.set(sqlSource, proxyNode);


        }


        return invocation.proceed();

    }

    private SqlNode proxyNode(SqlNode sqlnode) {
        SqlNode proxyNode = (SqlNode) Proxy.newProxyInstance(sqlnode.getClass().getClassLoader(),
                new Class[]{SqlNode.class}, new SqlNodeInvocationHandler(sqlnode));
        return proxyNode;
    }



    private class SqlNodeInvocationHandler implements InvocationHandler {

        private SqlNode target;

        public SqlNodeInvocationHandler(SqlNode target) {
            super();
            this.target = target;
        }

        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            DynamicContext context = (DynamicContext) args[0];
            context.getBindings().put(SCHEMA, schema);
            return method.invoke(target, args);
        }

    }

    public void setSchema(String schema) {
        if (StringUtils.isNotBlank(schema)) {
            if (!schema.endsWith(".")) {
                schema += ".";
            }
        }
        this.schema = schema;
    }

}
